import os
import tempfile
from seal import *

class KeyManager:
    def __init__(self, oss_handler, key_dir='keys'):
        self.oss_handler = oss_handler
        self.key_dir = key_dir
        os.makedirs(key_dir, exist_ok=True)

    def generate_keys(self, context):
        keygen = KeyGenerator(context)
        secret_key = keygen.secret_key()
        public_key = PublicKey()
        keygen.create_public_key(public_key)
        return secret_key, public_key

    def save_key(self, key, filename):
        file_path = os.path.join(self.key_dir, filename)
        key.save(file_path)
        return file_path

    def upload_private_key(self, key_path, oss_key):
        return self.oss_handler.upload_file(key_path, oss_key)

    def download_public_key(self, oss_key, local_path):
        return self.oss_handler.download_file(oss_key, local_path)

    def download_private_key(self, oss_key, local_path):
        return self.oss_handler.download_file(oss_key, local_path)
    def serialize_key(self, key):
        import pickle
        return pickle.dumps(key)
    def serialize_key_to_bytes(self, key):
        temp_file = tempfile.NamedTemporaryFile(delete=False)  # 不自动删除
        try:
            temp_file.close()  # 必须关闭，否则 SEAL 无法写入
            key.save(temp_file.name)

            with open(temp_file.name, 'rb') as f:
                key_bytes = f.read()
            return key_bytes
        except Exception as e:
            print(f"Serialization error: {e}")
            return None
        finally:
            try:
                os.remove(temp_file.name)  # 手动清理
            except Exception as cleanup_err:
                print(f"Cleanup warning: {cleanup_err}")